#include "EnhanceWrapper.h"

EnhanceWrapper::EnhanceWrapper(QString xmlPath, QObject *parent)
    : QObject{parent}
{
    qDebug() << ov::get_openvino_version().buildNumber << " " << ov::get_openvino_version().description;
    qDebug() << "Start To Compile CPU Model";
    ov::Core core;
    this->xmlModel = core.read_model(xmlPath.toStdString());
    printInputAndOutputsInfo(*this->xmlModel);
    double tS, tE;
    tS = (double) clock();
    this->compiledModel = core.compile_model(this->xmlModel, "CPU");
    tE = (double) clock();
    qDebug() << "Compile Done, cost " << (tE - tS) / (CLOCKS_PER_SEC) << " s";
}

cv::Mat EnhanceWrapper::infer(QString imageSrc){
    double tS, tE;

    tS = (double) clock();
    std::pair<cv::Mat, cv::Size> transPair = this->preprocess(imageSrc);
    tE = (double) clock();
    qDebug() << "Preprocess Done, cost " << (tE - tS) / (CLOCKS_PER_SEC) << " s";

    tS = (double) clock();
    QVector<float> buffer = this->_infer(transPair.first);
    tE = (double) clock();
    qDebug() << "Infer Done, cost " << (tE - tS) / (CLOCKS_PER_SEC) << " s";

    tS = (double) clock();
    cv::Mat res = this->postProcess(buffer, transPair.second);
    tE = (double) clock();
    qDebug() << "Postprocess Done, cost " << (tE - tS) / (CLOCKS_PER_SEC) << " s";

    return res;
}

cv::Mat EnhanceWrapper::postProcess(const QVector<float> &buffer, cv::Size oriSize){
    /*
        1. (np.transpose(input_image, (1, 2, 0)) + 1) / 2.0 * 255.0
        2. rgb -> bgr
        3. resize
    */
    std::vector<float> outBuffer(buffer.size());
    this->toHWC(buffer, outBuffer, cv::Size(512, 512));

    cv::Mat floatMat(512, 512, CV_32FC3, outBuffer.data());
    floatMat.convertTo(floatMat, CV_8UC3);
    cv::resize(floatMat, floatMat, oriSize, cv::INTER_CUBIC);
    return floatMat;
}

void EnhanceWrapper::toHWC(const QVector<float> &buffer, std::vector<float> &outBuffer, cv::Size bufferSize){
    const size_t dotNumber = bufferSize.width * bufferSize.height;
    const float scale = 1 / 2.0 * 255.0;
    for(size_t i = 0, j = 0; i < dotNumber * 3; j++){
        // b, g, r
        outBuffer[i] = (buffer[j + 2 * dotNumber] + 1)* scale;
        outBuffer[i + 1] = (buffer[j + dotNumber] + 1)* scale;
        outBuffer[i + 2] = (buffer[j] + 1) * scale;
        i += 3;
    }
}

QVector<float> EnhanceWrapper::_infer(cv::Mat transDim){
    ov::InferRequest inferRequest = this->compiledModel.create_infer_request();
    ov::Shape inputShape = {1, 3, 512, 512};
    ov::Tensor inputTensor = ov::Tensor(ov::element::f32, inputShape, transDim.data);
    inferRequest.set_input_tensor(inputTensor);

    // async
    inferRequest.start_async();
    inferRequest.wait();

    const ov::Tensor& outputTensor = inferRequest.get_output_tensor();
    const float *outputBuffer = outputTensor.data<const float>();
    return QVector<float>(outputBuffer, (outputBuffer + outputTensor.get_size()));
}

std::pair<cv::Mat, cv::Size> EnhanceWrapper::preprocess(QString imageSrc){
    return NetworkTool::ToolImagepreprocess(imageSrc);
}
